#!/usr/bin/env python
# -*-coding:utf-8 -*-
# @Time    : 2025-01-11 14:12:52
# @Author  : Huang Zeqing
# @File    : utils.py
# @Software: VS Code
# @contact : hzgreen6@gmail.com
# Description: 
#%%

from typing import Optional
from pathlib import Path
import pandas as pd
import re
import xarray as xr
import cfgrib
import numpy as np
from cfgrib import cfmessage
#%%

def delete_reftime(ds):

    use_coord = ['time', 'lat', 'lon']
    pattern = r'\.(\d{4})(\d{2})(\d{2})(\d{2})\.f'

    for coord in ds.coords:
        if coord not in use_coord:
            ds = ds.drop_vars(coord, errors='ignore')

    match = re.search(pattern, ds.encoding['source'])
    date_index = pd.to_datetime(match.group(1)+match.group(2)+match.group(3), format='%Y%m%d')
    ds = ds.assign_coords(time=pd.DatetimeIndex([date_index]))

    return ds


# function to create str for all dates from 2021 to 2024, with specified month, format like ""20210101"
def date_to_str(month):
    dates = []
    for year in range(2021, 2022): # 特定年份(更改)
        # Get the first and last day of the month
        start = pd.Timestamp(year=year, month=month, day=1)
        end = start + pd.offsets.MonthEnd(0)
        for day in pd.date_range(start, end, freq='D'):
            dates.append(day.strftime('%Y%m%d'))
    return dates

def encodings(vars: list) -> dict:
    """
    return a dict for saving nc file
    Args:
        vars: list of names of variables

    Returns:

    """

    return dict.fromkeys(vars, dict(zlib = True))

def make_path(path_list: list, parent_dir: Optional[Path] = None) -> Path:

    if parent_dir is not None:
        path = parent_dir.joinpath(*path_list)
    else:
        path = Path(*path_list)

    path.mkdir(parents = True, exist_ok=True)

    return path

def format_latlon_title(lat, lon):
    """
    format the title of lat and lon
    Args:
        lat: latitude
        lon: longitude

    Returns:

    """

    if lat >= 0:
        lat_title = f'{lat}°N'
    else:
        lat_title = f'{-lat}°S'

    if lon >= 0:
        lon_title = f'{lon}°E'
    else:
        lon_title = f'{-lon}°W'

    return f'{lat_title}, {lon_title}'

def with_open_xr(path, vars = None):
    """
    open a netCDF file with xarray
    Args:
    """

    import xarray as xr

    if vars is None:
        with xr.open_dataset(path) as ds:
            return ds
    else:
        with xr.open_dataset(path) as ds:
            return ds[vars]


def grb2_to_nc(grb2_path: str, nc_path: str, var_names: Optional[list] = None) -> None:
    """
    将GRIB2文件转换为NetCDF格式
    
    Args:
        grb2_path: GRIB2文件路径
        nc_path: 输出的NetCDF文件路径
        var_names: 需要保存的变量名列表，如果为None则保存所有变量
    """
    try:
        # 直接使用xarray读取GRIB2文件
        ds = xr.open_dataset(grb2_path, engine='cfgrib', backend_kwargs={'read_keys': ['typeOfLevel', 'level', 'stepRange']})
        
        # 如果指定了变量名，则只保留指定的变量
        if var_names is not None:
            ds = ds[var_names]
        
        # 设置压缩编码
        encoding = encodings(list(ds.data_vars))
        
        # 保存为NetCDF文件
        ds.to_netcdf(nc_path, encoding=encoding)
        
    except Exception as e:
        print(f"转换失败：{str(e)}")
    finally:
        # 确保数据集被正确关闭
        try:
            ds.close()
        except:
            pass



def nc_to_grb2(nc_path: str, grb2_path: str, var_names: Optional[list] = None) -> None:
    """
    将NetCDF文件转换为GRIB2格式
    
    Args:
        nc_path: NetCDF文件路径
        grb2_path: 输出的GRIB2文件路径
        var_names: 需要转换的变量名列表，如果为None则转换所有变量
    """
    try:
        # 读取NetCDF文件
        ds = xr.open_dataset(nc_path)
        
        # 如果指定了变量名，则只保留指定的变量
        if var_names is not None:
            ds = ds[var_names]
        
        # 设置GRIB2必需的参数
        grib_kwargs = {
            'edition': 2,
            'centre': 'kwbc',  # NCEP中心
            'gridType': 'regular_ll',  # 经纬度网格
            'typeOfLevel': 'surface',
            'dataDate': int(ds.time.dt.strftime('%Y%m%d').values[0]),
            'dataTime': int(ds.time.dt.strftime('%H%M').values[0])
        }
        
        # 转换每个变量并写入GRIB2文件
        for var_name in ds.data_vars:
            var_data = ds[var_name].values
            message = cfmessage.CfMessage({
                'values': var_data.flatten(),
                'paramId': 1,  # 参数ID，需要根据实际变量类型设置
                'shortName': var_name,
                'gridType': 'regular_ll',
                'Nx': var_data.shape[-1],
                'Ny': var_data.shape[-2],
                'distinctLatitudes': ds.lat.values,
                'distinctLongitudes': ds.lon.values,
                **grib_kwargs
            })
            message.write(grb2_path)
        
        print(f"转换完成：{nc_path} -> {grb2_path}")
        
    except Exception as e:
        print(f"转换失败：{str(e)}")
    finally:
        # 确保数据集被正确关闭
        try:
            ds.close()
        except:
            pass

